# -*- coding: utf-8 -*-
import cv2
import numpy as np
import tqdm
import json
import os
from pycocotools.coco import COCO
import shutil
import matplotlib.pyplot as plt

jsonPath = "./train_a.json"
imgPath = "./train_a_imgs/"
colorImgPath = "../imgs/"
binaryImgPath = "../masks/"

coco = COCO(jsonPath)
imgIds = coco.getImgIds()
for i in imgIds:
    img = coco.loadImgs(imgIds[i])[0]
    shutil.copy(os.path.join(imgPath, img['file_name']), colorImgPath)
    img_bmp = cv2.imread(os.path.join(imgPath, img['file_name']))
    img_name = img['file_name'].split(".")[0]
    cv2.imwrite(colorImgPath + img_name + ".png", img_bmp)
    catIds = []
    for ann in coco.dataset['annotations']:
        if ann['image_id'] == imgIds[i]:
            catIds.append(ann['category_id'])
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    width = img['width']
    height = img['height']
    anns = coco.loadAnns(annIds)
    # mask_pic = np.zeros((height, width))
    imgs = np.zeros(shape=(int(height), int(width)), dtype=np.float32)
    for single in anns:
        mask_single = coco.annToMask(single)
        # cv2.imshow('a', mask_single)
        # cv2.waitKey(0)
        if single['category_id'] != 0:
            for row in range(height):
                for col in range(width):
                    if mask_single[row][col] > 0:
                        imgs[row][col] = 1
    imgs = imgs.astype(np.uint8)
    img_name = img['file_name'].split(".")[0]
    cv2.imwrite(binaryImgPath + img_name + ".png", imgs)
